{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import shutil\n",
    "import pathlib\n",
    "\n",
    "import time_machine\n",
    "\n",
    "# import to register magics\n",
    "import sqlmesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e38a8ec8-3c54-4a21-aa28-a6fd1fdc86e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "freezer = time_machine.travel(\"2032-01-01 00:00:00 UTC\")\n",
    "freezer.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a412470c028d252",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "starting_dir = pathlib.Path(os.getcwd())\n",
    "prev_parent_stem = \"\"\n",
    "root_dir = None\n",
    "for parent in starting_dir.parents:\n",
    "    if parent.stem == \"sqlmesh\" and prev_parent_stem == \"tests\":\n",
    "        root_dir = parent\n",
    "        break\n",
    "    prev_parent_stem = parent.stem\n",
    "else:\n",
    "    raise RuntimeError(\"Couldn't find root dir\")\n",
    "\n",
    "EXAMPLE_SUSHI_DIR = pathlib.Path(root_dir) / \"examples\" / \"sushi\"\n",
    "str(EXAMPLE_SUSHI_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78afd8b3-3f88-4bfc-a3e7-52770c64a647",
   "metadata": {},
   "outputs": [],
   "source": [
    "TARGET_SUSHI_DIR = pathlib.Path(tempfile.gettempdir()) / \"test_jupyter_notebook\"\n",
    "str(TARGET_SUSHI_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e36b791-d756-4a01-83f5-d126ca69621c",
   "metadata": {},
   "outputs": [],
   "source": [
    "TARGET_MODEL_DIR = TARGET_SUSHI_DIR / \"models\"\n",
    "str(TARGET_MODEL_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0f43c20-7fe4-448b-abc8-ced4dcb87ad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ignore(src, names):\n",
    "    if pathlib.Path(src).name in {\".cache\", \"__pycache__\", \"logs\", \"data\"}:\n",
    "        return names\n",
    "    return []\n",
    "\n",
    "if TARGET_SUSHI_DIR.exists():\n",
    "    shutil.rmtree(str(TARGET_SUSHI_DIR))\n",
    "shutil.copytree(str(EXAMPLE_SUSHI_DIR), str(TARGET_SUSHI_DIR), ignore=ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6121c12f2cf70f05",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%context $TARGET_SUSHI_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0352bd04-bdef-476e-9b88-ca9f1a6a4e59",
   "metadata": {},
   "source": [
    "### Leaving plan unapplied to verify that the text boxes are printed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1055eaa-533d-4bf5-adbc-4fb5b5e42aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%plan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825596d3292f4cf4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%plan --no-prompts --auto-apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4478d80a-ee61-43b5-ac48-60b20a497a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "%plan dev --no-prompts --auto-apply --include-unmodified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "495225e4-1e23-4a76-87a2-881d1bb3b017",
   "metadata": {},
   "outputs": [],
   "source": [
    "%render sushi.top_waiters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e40f213-5040-4659-957e-7fd478ce4c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "%render sushi.top_waiters --no-format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb381515-c214-4b85-83fd-136a1c2fb525",
   "metadata": {},
   "outputs": [],
   "source": [
    "%evaluate sushi.top_waiters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9241c2a5-1ca6-42db-9219-b7fd9ae0f118",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_waiters = TARGET_MODEL_DIR / \"top_waiters.sql\"\n",
    "top_waiters.write_text(\"\"\"/* View of top waiters. */\n",
    "MODEL (\n",
    "  name sushi.top_waiters,\n",
    "  owner jen,\n",
    "  audits (\n",
    "    unique_values(columns=[waiter_id])\n",
    "  ),\n",
    "  grain waiter_id\n",
    ");\n",
    "\n",
    "SELECT\n",
    "  waiter_id::INT AS waiter_id,\n",
    "  revenue::DOUBLE AS revenue,\n",
    "  1 as new_col\n",
    "FROM sushi.waiter_revenue_by_day\n",
    "WHERE\n",
    "  event_date = (\n",
    "    SELECT\n",
    "      MAX(event_date)\n",
    "    FROM sushi.waiter_revenue_by_day\n",
    "  )\n",
    "ORDER BY\n",
    "  revenue DESC\n",
    "LIMIT 10\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9cc893-5954-4195-a0bc-785244678e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "%diff prod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7ffb02-4ab1-48ca-949c-908f0129c0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%plan dev --no-prompts --auto-apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e6fb605-f6fd-469b-a02c-cbd5d7f662ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%dag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5daa9a5d-93d7-489c-b9ae-fc62039c3208",
   "metadata": {},
   "outputs": [],
   "source": [
    "%migrate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f35db811-3874-4faf-b74e-e6cae3dbd1b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%table_diff prod:dev --model sushi.top_waiters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a3d152-19b6-47a0-8ca2-fe95cb5484c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%run_test {TARGET_SUSHI_DIR}/tests/test_customer_revenue_by_day.yaml::test_customer_revenue_by_day"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad5ed56-89cd-4eb0-a64c-ab6d26e92e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%audit sushi.top_waiters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "347cbe5d-e8d4-4d0f-b9c9-9b920c8df314",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%fetchdf my_result\n",
    "SELECT 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a353df-26c0-4ab6-91ec-c8bab0847c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eaf0519-bba8-4896-8c0f-1b30d9777988",
   "metadata": {},
   "outputs": [],
   "source": [
    "%info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db0cc0ee-fd63-4376-8562-d5f9eab43710",
   "metadata": {},
   "outputs": [],
   "source": [
    "%migrate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c5009e8-fa98-48de-8d40-ec274d6cd97f",
   "metadata": {},
   "outputs": [],
   "source": [
    "freezer = time_machine.travel(\"2032-01-02 00:00:00\")\n",
    "freezer.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f35ae832-de18-4d91-96f9-955af2cd9198",
   "metadata": {},
   "outputs": [],
   "source": [
    "%run_dag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0138a5df-f136-47fd-8146-62c762ef776d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%invalidate dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "483ca2f5-585b-4baf-8a0a-77f6aac5c31e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "UsageError: Line magic function `%context` not found.\n"
     ]
    }
   ],
   "source": [
    "%context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "04d73a18-58d3-4599-9d56-7767c4320be1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlmesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "38116170-c40c-45cb-88d1-f6198020ec76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "::\n",
       "\n",
       "  %create_test --query QUERY [QUERY ...] [--overwrite]\n",
       "                   [--var VAR [VAR ...]] [--path PATH] [--name NAME]\n",
       "                   model\n",
       "\n",
       "Generate a unit test fixture for a given model.\n",
       "\n",
       "positional arguments:\n",
       "  model\n",
       "\n",
       "options:\n",
       "  --query <QUERY [QUERY ...]>, -q <QUERY [QUERY ...]>\n",
       "                        Queries that will be used to generate data for the\n",
       "                        model's dependencies.\n",
       "  --overwrite, -o       When true, the fixture file will be overwritten in\n",
       "                        case it already exists.\n",
       "  --var <VAR [VAR ...]>, -v <VAR [VAR ...]>\n",
       "                        Key-value pairs that will define variables needed by\n",
       "                        the model.\n",
       "  --path PATH, -p PATH  The file path corresponding to the fixture, relative\n",
       "                        to the test directory. By default, the fixture will be\n",
       "                        created under the test directory and the file name\n",
       "                        will be inferred based on the test's name.\n",
       "  --name NAME, -n NAME  The name of the test that will be created. By default,\n",
       "                        it's inferred based on the model's name.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/repos/sqlmesh/sqlmesh/magics.py"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%create_test?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "849f85f0-f116-4a33-865f-e00f7176566f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mDocstring:\u001b[0m List currently available magic functions.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/.pyenv/versions/3.11.5/envs/sqlmesh311/lib/python3.11/site-packages/IPython/core/magics/basic.py"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%lsmagic?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e85f5c7-37a9-4b32-8c76-ee444a83a36b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
